import logging
import time
from contextlib import contextmanager, AbstractContextManager
from dataclasses import dataclass
from multiprocessing import Process, Queue, Manager
from typing import Dict, Literal, List

from storage._perf import profiling
from storage.kvstore.base import AbstractKVStore


class BackgroundWriter(Process):
    @dataclass
    class Intent:
        action: Literal['put', 'delete']
        keys: List[bytes] = None
        pairs: Dict[bytes, bytes] = None

    def __init__(self, q: Queue, store: AbstractKVStore, pending_data: Dict[bytes, bytes]):
        super().__init__(name='background-writer')
        self._queue = q
        self._store = store
        self._pending = pending_data
        self._logger = logging.getLogger('fingerprint.background-writer')

    def run(self):
        while item := self._queue.get():
            if item.action == 'put':
                self._store.update(item.pairs)
                for k in item.pairs:
                    del self._pending[k]
            elif item.action == 'delete':
                for k in item.keys:
                    del self._store[k]


class MultiProcessKVStore(AbstractKVStore):
    def __init__(self, store: AbstractKVStore, *, max_size=1000):
        self._queue = Queue(max_size)
        self._store = store
        self._mgr = Manager()
        self._pending = self._mgr.dict()
        self._process = BackgroundWriter(self._queue, self._store, self._pending)

    def close(self, *, graceful_wait: int = 10):
        elapsed = 0
        self._process.terminate()
        while self._process.is_alive():
            start = time.time()
            time.sleep(1.0)
            elapsed = time.time() - start
            if elapsed > graceful_wait:
                self._process.kill()

        if hasattr(self._store, 'close'):
            self._store.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get(self, k: bytes, default: bytes = None) -> bytes:
        if ret := self._pending.get(k):
            return ret
        return self._store.get(k, default)

    def put(self, k: bytes, v: bytes):
        self._pending[k] = v
        self._queue.put(BackgroundWriter.Intent('put', pairs={k: v}))

    @profiling(output='./lmdb_update.pstats', interval_call=1000)
    def update(self, kv_pairs: Dict[bytes, bytes]):
        for k, v in kv_pairs.items():
            self._pending[k] = v

        self._queue.put(BackgroundWriter.Intent('put', pairs=kv_pairs))

    def __delitem__(self, k: bytes):
        del self._pending[k]
        self._queue.put(BackgroundWriter.Intent('delete', keys=[k]))

    def __getitem__(self, item: bytes):
        data = self.get(item)
        if data is None:
            raise KeyError(item)

    def __setitem__(self, key: bytes, value: bytes):
        self.put(key,value)

    @contextmanager
    def tx(self) -> AbstractContextManager:
        with self._store.tx():
            yield
